import cv2
from super_gradients.training import models
import torch

# 加载 YOLO-NAS Nano 模型
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = models.get('yolo_nas_n', pretrained_weights="coco").to(device)

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    # 读取一帧图像
    ret, frame = cap.read()
    if not ret:
        break

    # 调整输入图像尺寸，减小计算量
    resized_frame = cv2.resize(frame, (320, 320))

    # 使用模型进行目标检测
    detections = list(model.predict(resized_frame, conf=0.3))[0]

    # 获取检测结果
    bboxes_xyxy = detections.prediction.bboxes_xyxy.cpu().numpy().astype(int)
    confidence = detections.prediction.confidence.cpu().numpy()
    labels = detections.prediction.labels.cpu().numpy().astype(int)

    # 绘制检测框和标签
    for (bbox, conf, label) in zip(bboxes_xyxy, confidence, labels):
        x1, y1, x2, y2 = bbox
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, f'{label}: {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 显示带检测结果的图像
    cv2.imshow('YOLO-NAS Nano Detection', frame)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭所有窗口
cap.release()
cv2.destroyAllWindows()
    